import numpy as np
np.int = int  # 强制覆盖np.int的指向
from sko.GA import GA_TSP
from sko.SA import SA_TSP
from sko.ACA import ACA_TSP
import matplotlib.pyplot as plt

# ================== 全局参数配置 ==================
CITY_NUM = 50  # 城市数量
POP_SIZE = 400  # 种群规模
MAX_ITER = 500  # 迭代次数
MUTATE_PROB = 0.35  # 变异概率(遗传算法)
T0 = 100  # 初始温度(模拟退火)
Q = 0.98  # 温度衰减系数(模拟退火)


# ================== 核心功能类 ==================
class TSPSolver:
    def __init__(self, city_num=CITY_NUM):
        self.cities = self.generate_cities(city_num)
        self.distance_matrix = self.cal_distance_matrix()

    @staticmethod
    def generate_cities(n):
        """生成随机城市坐标"""
        np.random.seed(100)
        return np.random.rand(n, 2) * 100

    def cal_distance_matrix(self):
        """计算距离矩阵"""
        n = len(self.cities)
        dist_mat = np.zeros((n, n))
        for i in range(n):
            for j in range(i + 1, n):
                dist = np.linalg.norm(self.cities[i] - self.cities[j])
                dist_mat[i][j] = dist_mat[j][i] = dist
        return dist_mat

    def path_distance(self, path):
        """计算路径总距离"""
        total = 0
        for i in range(len(path) - 1):
            total += self.distance_matrix[path[i], path[i + 1]]
        total += self.distance_matrix[path[-1], path[0]]  # 回到起点
        return total

    # ========== 算法实现部分 ==========
    def ga_solve(self, n_pop=POP_SIZE, max_iter=MAX_ITER, mutation_prob=MUTATE_PROB):
        """遗传算法求解"""
        ga = GA_TSP(
            func=self.path_distance,
            n_dim=CITY_NUM,
            size_pop=n_pop,
            max_iter=max_iter,
            prob_mut=mutation_prob
        )
        best_path, best_distance = ga.run()
        # 关键修正：确保best_distance为标量
        if isinstance(best_distance, np.ndarray):
            best_distance = best_distance[0]
        return best_path, best_distance, ga.generation_best_Y

    def sa_solve(self, T0=T0, Q=Q, max_iter=MAX_ITER):
        """模拟退火算法求解"""
        x0 = np.random.permutation(CITY_NUM)
        sa = SA_TSP(
            func=self.path_distance,
            n_dim=CITY_NUM,
            x0=x0,
            T0=T0,
            T_min=1e-3,
            L=max_iter // 10,
            max_stay_counter=max_iter // 10,
            Q=Q
        )
        best_path, best_distance = sa.run()
        return best_path, best_distance, sa.best_y_history

    def aca_solve(self, n_pop=POP_SIZE, max_iter=MAX_ITER):
        """蚁群算法求解"""
        aca = ACA_TSP(
            func=self.path_distance,
            n_dim=CITY_NUM,
            size_pop=n_pop,
            max_iter=max_iter,
            alpha=1.0,
            beta=2.0,
            distance_matrix=self.distance_matrix
        )
        best_path, best_distance = aca.run()
        return best_path, best_distance, aca.y_best_history

    # ========== 可视化方法 ==========
    @staticmethod
    def plot_convergence(data_dict, title='Algorithm Comparison'):
        """绘制收敛曲线对比"""
        plt.figure(figsize=(10, 6))
        for name, data in data_dict.items():
            plt.plot(data, label=name)
        plt.title(title)
        plt.xlabel('Iteration')
        plt.ylabel('Best Distance')
        plt.legend()
        plt.grid(True)
        plt.show()

    def plot_path(self, path, title='Optimal Path'):
        """绘制路径图"""
        plt.figure(figsize=(10, 6))
        plt.scatter(self.cities[:, 0], self.cities[:, 1], c='r', marker='o')

        for i in range(len(path)):
            start = self.cities[path[i]]
            end = self.cities[path[(i + 1) % len(path)]]
            plt.plot([start[0], end[0]], [start[1], end[1]], 'b-')

        plt.title(f'{title} (Distance: {self.path_distance(path):.2f})')
        plt.xlabel('X Coordinate')
        plt.ylabel('Y Coordinate')
        plt.show()


# ================== 主程序 ==================
if __name__ == '__main__':
    solver = TSPSolver(CITY_NUM)

    # 运行不同算法
    ga_path, ga_dist, ga_hist = solver.ga_solve()
    sa_path, sa_dist, sa_hist = solver.sa_solve()
    aca_path, aca_dist, aca_hist = solver.aca_solve()

    # 结果对比
    print(f"Genetic Algorithm: {ga_dist:.2f}")
    print(f"Simulated Annealing: {sa_dist:.2f}")
    print(f"Ant Colony: {aca_dist:.2f}")

    # 绘制收敛曲线对比
    solver.plot_convergence({
        'Genetic Algorithm': ga_hist,
        'Simulated Annealing': sa_hist[:len(ga_hist)],
        'Ant Colony': aca_hist[:len(ga_hist)]
    })

    # 绘制各算法最优路径
    solver.plot_path(ga_path, 'GA Optimal Path')
    solver.plot_path(sa_path, 'SA Optimal Path')
    solver.plot_path(aca_path, 'ACA Optimal Path')